import random
import argparse

def read_file(file_path):
    """读取文件，返回包含照片路径的列表"""
    with open(file_path, 'r') as file:
        return file.readlines()

def split_data(file_paths, val_size=None, val_ratio=None):
    """根据指定比例或数量划分训练集和验证集"""
    # 打乱数据顺序
    random.shuffle(file_paths)

    # 根据数量或比例计算验证集大小
    if val_size:
        val_size = min(val_size, len(file_paths))
    elif val_ratio:
        val_size = int(len(file_paths) * val_ratio)
    else:
        raise ValueError("必须指定验证集的比例或数量")

    # 划分训练集和验证集
    val_set = file_paths[:val_size]
    train_set = file_paths[val_size:]

    return train_set, val_set

def write_to_file(data, file_path):
    """将数据写入文件"""
    with open(file_path, 'w') as file:
        file.writelines(data)

def main(input_txt, train_txt, val_txt, val_size=None, val_ratio=None):
    # 读取txt文件中的照片路径
    file_paths = read_file(input_txt)

    # 打印原始数据集数量
    print(f"原始数据集数量: {len(file_paths)}")

    # 切分数据为训练集和验证集
    train_set, val_set = split_data(file_paths, val_size, val_ratio)

    # 打印训练集和验证集数量
    print(f"训练集数量: {len(train_set)}")
    print(f"验证集数量: {len(val_set)}")

    # 写入结果到文件
    write_to_file(train_set, train_txt)
    write_to_file(val_set, val_txt)

    print(f"训练集保存到: {train_txt}")
    print(f"验证集保存到: {val_txt}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="划分数据为训练集和验证集")
    parser.add_argument('input_txt', type=str, help="包含照片路径的txt文件")
    parser.add_argument('train_txt', type=str, help="训练集输出txt文件")
    parser.add_argument('val_txt', type=str, help="验证集输出txt文件")
    parser.add_argument('--val_size', type=int, default=None, help="验证集数量")
    parser.add_argument('--val_ratio', type=float, default=None, help="验证集比例")

    args = parser.parse_args()

    main(args.input_txt, args.train_txt, args.val_txt, args.val_size, args.val_ratio)

# python txt_split.py gt_guiji_all_20240909_subset3.txt gt_guiji_all_20240909_subset3_train.txt gt_guiji_all_20240909_subset3_val.txt --val_ratio 0.1
# python txt_split.py gt_guiji_all_20240909_subset3.txt gt_guiji_all_20240909_subset3_train.txt gt_guiji_all_20240909_subset3_val.txt --val_size 50000
